from __future__ import absolute_import, print_function
import torch
import torchvision
import torch.nn as nn
from torch.nn import functional as F
from IPython import embed


class PersonResNet(nn.Module):
    def __init__(self, num_classes, pretrained=True, loss={'softmax, metric'},  **kwargs):
        super(PersonResNet, self).__init__()
        self.pretrained = pretrained
        for k, v in kwargs.items():
            self.__setattr__(k, v)

        # resnetN = torchvision.models.resnet18(pretrained=True)
        resnetN = torchvision.models.resnet50(pretrained=self.pretrained)
        # resnetN = torchvision.models.resnet101(pretrained=True)
        # resnetN = torchvision.models.resnet152(pretrained=True)

        in_feat = resnetN.fc.in_features
        # print(in_feat)

        # print(resnetN)
        # for i in list(resnetN.children()):
        #     print(i)
        #     print(50*'=')
        # embed()

        # 特征提取器
        # 想要换掉最后的(avgpool)和(fc)
        # 取出前面的子模块，解包，组成Sequential
        self.base = nn.Sequential(*list(resnetN.children())[:-2])

        # 汇集器，在forward里使用nn.functional来实现

        # 分类器
        self.classifier = nn.Linear(in_feat, num_classes)

    def forward(self, x, training=True):
        # 特征提取
        x = self.base(x)

        # 汇集 --> [B x C x 1 x 1]
        x = F.avg_pool2d(x, x.shape[-2:])

        # 展平 4D -> 2D 得到 reid-feature
        # feat_ = x.view(x.shape[:2])  # 法一，按原 [B, C] 排布
        feat_ = x.view(x.shape[0], -1)  # 法二，保留B，自动推断
        # feat_ = x.squeeze(dim=-1).squeeze(dim=-1)  # 法三，压缩后两维

        # print("shape of reid-feature:", feat_.shape)

        # ============= 特征向量的归一化 ==========================
        feat = 1. * feat_ / (torch.norm(feat_, 2, dim=-1,
                             keepdim=True).expand_as(feat_) + 1e-12)
        # =================================================================

        # 测试阶段只需要 f 做度量
        if not training:
            return feat_, feat

        # 分类
        y = self.classifier(feat)

        # 训练阶段需要 f做度量学习 和 y做表征学习
        return feat_, feat, y


if __name__ == '__main__':
    model = PersonResNet(751)
    x = torch.ones((2, 3, 384, 128), dtype=torch.float32)

    feat_, feat, y = model(x)
    print("feat_:", feat_.shape, "feat:", feat.shape, "y:", y.shape)

    feat_, feat = model(x, training=False)
    print("feat_:", feat_.shape, "feat:", feat.shape)
